import pandas as pd
from matplotlib import pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller
from statsmodels.tsa.seasonal import seasonal_decompose
import numpy as np

# 读取数据
data = pd.read_csv("D:/五粮液股票数据.csv")
# 将数据处理为时间序列数据
datas = pd.to_datetime(data["交易日期"].values, format="%Y%m%d")
# 将时间序列设置为数据的索引
shares_info = data.set_index(datas)
# 使用收盘价分析股票价格，按周重采样，取出每周的均值
shares_week = shares_info["收盘价"].resample(rule="W-MON").mean()
# 截取部分数据作为训练集
ts = shares_week["2014":"2017"]


# 作移动平均图，观察消除可能存在的季节性后图像的样子
def draw_trend(timeSeries, size):
    f = plt.figure(facecolor='white')
    # 对size个数据进行移动平均
    rol_mean = timeSeries.rolling(window=size).mean()
    # 对size个数据进行加权移动平均
    rol_weighted_mean = timeSeries.ewm(span=size).mean()

    timeSeries.plot(color='blue', label='Original')
    rol_mean.plot(color='red', label='Rolling Mean')
    rol_weighted_mean.plot(color='black', label='Weighted Rolling Mean')
    plt.legend(loc='best')
    plt.title('Rolling Mean')
    plt.show()


# 绘制timeSeries对象
def draw_ts(timeSeries):
    f = plt.figure(facecolor='white')
    timeSeries.plot(color='blue')
    plt.show()


# 单位根检验：
# 增广迪基-富勒（ADF）检验的原假设是存在单位根，备择假设是不存在单位根。换言之，p值越大，我们越有理由认为存在单位根，也即原时间序列越不平稳，不能用ARIMA模型进行拟合
def testStationarity(ts):
    dftest = adfuller(ts)
    # 对上述函数求得的值进行语义描述
    dfoutput = pd.Series(dftest[0:4], index=['Test Statistic', 'p-value', '#Lags Used', 'Number of Observations Used'])
    for key, value in dftest[4].items():
        dfoutput['Critical Value (%s)' % key] = value
    if dftest[0] > dftest[4]["10%"]:
        print("数据平稳性差")
    elif dftest[0] > dftest[4]["5%"]:
        print("有90%以上的b把握认为数据平稳性好")
    elif dftest[0] > dftest[4]["1%"]:
        print("有99%以上的b把握认为数据平稳性好")
    else:
        print("数据平稳性好")
    print(dfoutput)
    return dfoutput


# 自相关和偏相关图，默认阶数为31阶
def draw_acf_pacf(ts, lags=31):
    f = plt.figure(facecolor='white')
    ax1 = f.add_subplot(211)
    plot_acf(ts, lags=31, ax=ax1)
    ax2 = f.add_subplot(212)
    plot_pacf(ts, lags=31, ax=ax2)
    plt.show()


# 以下为四种提高数据平稳性的手段：
# # 进行取对数处理，降低原数据的波动幅度
# ts_log = np.log(ts)
# draw_ts(ts_log)
#
# # 观察原始数据的取移动平均和指数平均后的整体趋势（平滑法）。对于具有明显季节性规律的数据，进行平滑处理可以一定程度上消减季节差异
# draw_trend(ts_log, 12)
#
# # 进行差分。差分是十分重要的平稳化手段，一般一到两次差分足以使得原始数据变平稳。差分次数过多则会使得数据变成完全的随机序列
# diff_12 = ts_log.diff(12)
# diff_12.dropna(inplace=True)  # 删去因差分而空出的数据（如原长m的列表进行差分后只有m-1个有效值，剩下那一个是NaN
# diff_12_1 = diff_12.diff(1)
# diff_12_1.dropna(inplace=True)
# testStationarity(diff_12_1)
#
# # 将数据进行加法分解（如果想用乘法算法可以把实参改为"multiplicative"，分别观察原数据具有长期特性、季节趋势和完全随机的分量
# decomposition = seasonal_decompose(ts_log, model="additive")
# trend = decomposition.trend
# seasonal = decomposition.seasonal
# residual = decomposition.resid

# 通过观察，数据波动性较大，有明显的长期特性，但季节趋势不突出，因此对它取对数和差分，消减这两个性质带给它的不平稳性。移动平均就不进行了
ts_log = np.log(ts)
# rol_mean = ts_log.rolling(window=12).mean()
# rol_mean.dropna(inplace=True)
ts_diff_1 = ts_log.diff(1)
ts_diff_1.dropna(inplace=True)
testStationarity(ts_diff_1)
# 如果一次差分不明显还可以再来一次。此处虽然已经十分平稳了，但直接进行拟合会无法收敛，因此还要再差分一次
ts_diff_2 = ts_diff_1.diff(1)
ts_diff_2.dropna(inplace=True)
testStationarity(ts_diff_1)

model = ARIMA(ts_diff_2, order=(1, 1, 0))  # 求解模型。有时会遇到AR、I、MA三模型中部分模型的参数无法反转而报错
# 只需在上行元组里把相应参数设置为0即可
result_arima = model.fit()

# 进行预测前要把求解出的模型还原
predict_ts = result_arima.predict()
# 一阶差分还原
diff_shift_ts = ts_diff_1.shift(1)
diff_recover_1 = predict_ts.add(diff_shift_ts)
# 再次一阶差分还原
rol_shift_ts = ts_log.shift(1)
diff_recover = diff_recover_1.add(rol_shift_ts)
# # 移动平均还原
# rol_sum = ts_log.rolling(window=11).sum()
# rol_recover = diff_recover*12 - rol_sum.shift(1)
# 对数还原
log_recover = np.exp(rol_shift_ts)
log_recover.dropna(inplace=True)

ts = ts[log_recover.index]  # 过滤没有预测的记录
plt.figure(facecolor='white')
log_recover.plot(color='blue', label='Predict')
ts.plot(color='red', label='Original')
plt.legend(loc='best')
plt.title('RMSE: %.4f' % np.sqrt(sum((log_recover-ts)**2)/ts.size))  # 使用均方根误差评判预测好坏
plt.show()
